use futures::{SinkExt, StreamExt};
use warp::ws::{Message, WebSocket};
use warp::Filter;

// 定义WebSocket的路由
pub fn websocket_route() -> impl Filter<Extract = impl warp::Reply, Error = warp::Rejection> + Clone {
    warp::path("ws")
        .and(warp::ws())
        .map(|ws: warp::ws::Ws| {
            ws.on_upgrade(handle_connection)
        })
}

// 处理WebSocket连接
async fn handle_connection(ws: WebSocket) {
    // 将WebSocket消息拆分为接收和发送部分
    let (mut ws_tx, mut ws_rx) = ws.split();

    // 处理接收到的每一个消息
    while let Some(result) = ws_rx.next().await {
        match result {
            Ok(msg) => {
                if msg.is_text() {
                    // 处理文本消息
                    let reply = format!("Received: {}", msg.to_str().unwrap());
                    ws_tx.send(Message::text(reply)).await.unwrap();
                } else if msg.is_binary() {
                    // 处理二进制消息
                    ws_tx.send(Message::binary(msg.as_bytes())).await.unwrap();
                }
            }
            Err(e) => {
                eprintln!("websocket error: {}", e);
                break;
            }
        }
    }
}
